#!/usr/bin/env python
"""
数据库备份脚本 - 创建数据库的备份并管理备份文件
支持MySQL/MariaDB、PostgreSQL和SQLite数据库
"""

import os
import sys
import time
import shutil
import logging
import subprocess
import django
from datetime import datetime
from pathlib import Path

# 添加项目根目录到Python路径
BASE_DIR = Path(__file__).resolve().parent.parent
sys.path.append(str(BASE_DIR))

# 设置Django环境
os.environ.setdefault('DJANGO_SETTINGS_MODULE', 'boxuser.settings')
django.setup()

# 导入Django设置
from django.conf import settings

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s [%(levelname)s] %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler(BASE_DIR / 'logs' / 'backup.log')
    ]
)
logger = logging.getLogger(__name__)

# 确保日志目录存在
os.makedirs(BASE_DIR / 'logs', exist_ok=True)

# 备份配置
BACKUP_DIR = BASE_DIR / 'backups'
MAX_BACKUPS = 10  # 保留的最大备份数量


def ensure_backup_dir():
    """确保备份目录存在"""
    os.makedirs(BACKUP_DIR, exist_ok=True)
    logger.info(f"备份目录: {BACKUP_DIR}")


def get_db_settings():
    """从Django设置中获取数据库配置"""
    db_settings = settings.DATABASES.get('default', {})
    
    if not db_settings:
        raise ValueError("未找到默认数据库配置")
    
    return {
        'engine': db_settings.get('ENGINE', ''),
        'name': db_settings.get('NAME', ''),
        'user': db_settings.get('USER', ''),
        'password': db_settings.get('PASSWORD', ''),
        'host': db_settings.get('HOST', ''),
        'port': db_settings.get('PORT', ''),
    }


def backup_mysql(db_settings):
    """备份MySQL/MariaDB数据库"""
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_file = BACKUP_DIR / f"backup_{timestamp}.sql"
    
    # 构建mysqldump命令
    cmd = [
        'mysqldump',
        '--single-transaction',
        '--quick',
        '--lock-tables=false',
        f"--user={db_settings['user']}",
    ]
    
    if db_settings['password']:
        cmd.append(f"--password={db_settings['password']}")
    
    if db_settings['host']:
        cmd.append(f"--host={db_settings['host']}")
    
    if db_settings['port']:
        cmd.append(f"--port={db_settings['port']}")
    
    cmd.append(db_settings['name'])
    
    try:
        logger.info(f"开始备份MySQL数据库 {db_settings['name']}...")
        with open(backup_file, 'w') as f:
            process = subprocess.run(
                cmd,
                stdout=f,
                stderr=subprocess.PIPE,
                text=True
            )
        
        if process.returncode != 0:
            logger.error(f"备份失败: {process.stderr}")
            return None
        
        # 检查备份文件是否为空或太小
        if os.path.getsize(backup_file) < 100:
            with open(backup_file, 'r') as f:
                content = f.read()
            logger.error(f"备份可能失败，文件过小或为空。内容: {content}")
            return None
        
        logger.info(f"备份成功: {backup_file}")
        return backup_file
    
    except Exception as e:
        logger.error(f"备份过程中发生错误: {e}")
        return None


def backup_postgresql(db_settings):
    """备份PostgreSQL数据库"""
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_file = BACKUP_DIR / f"backup_{timestamp}.sql"
    
    # 环境变量用于pg_dump认证
    env = os.environ.copy()
    if db_settings['password']:
        env['PGPASSWORD'] = db_settings['password']
    
    # 构建pg_dump命令
    cmd = [
        'pg_dump',
        '--dbname', db_settings['name'],
        '--username', db_settings['user'],
        '--format', 'custom',
        '--file', str(backup_file),
    ]
    
    if db_settings['host']:
        cmd.extend(['--host', db_settings['host']])
    
    if db_settings['port']:
        cmd.extend(['--port', db_settings['port']])
    
    try:
        logger.info(f"开始备份PostgreSQL数据库 {db_settings['name']}...")
        process = subprocess.run(
            cmd, 
            env=env,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True
        )
        
        if process.returncode != 0:
            logger.error(f"备份失败: {process.stderr}")
            return None
        
        logger.info(f"备份成功: {backup_file}")
        return backup_file
    
    except Exception as e:
        logger.error(f"备份过程中发生错误: {e}")
        return None


def backup_sqlite(db_settings):
    """备份SQLite数据库"""
    db_file = db_settings['name']
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    backup_file = BACKUP_DIR / f"backup_{timestamp}.sqlite3"
    
    try:
        logger.info(f"开始备份SQLite数据库 {db_file}...")
        # 检查文件是否存在
        if not os.path.exists(db_file):
            logger.error(f"数据库文件不存在: {db_file}")
            return None
            
        shutil.copy2(db_file, backup_file)
        logger.info(f"备份成功: {backup_file}")
        return backup_file
    
    except Exception as e:
        logger.error(f"备份过程中发生错误: {e}")
        return None


def cleanup_old_backups():
    """清理旧备份文件，只保留最近的几个"""
    backups = sorted(
        [f for f in BACKUP_DIR.glob('backup_*')],
        key=os.path.getmtime
    )
    
    # 如果备份数量超过最大限制，删除最旧的备份
    if len(backups) > MAX_BACKUPS:
        backups_to_delete = backups[:-MAX_BACKUPS]
        for backup in backups_to_delete:
            try:
                os.remove(backup)
                logger.info(f"已删除旧备份: {backup}")
            except Exception as e:
                logger.error(f"删除旧备份 {backup} 时出错: {e}")


def run_backup():
    """运行数据库备份流程"""
    logger.info("开始数据库备份过程...")
    
    # 确保备份目录存在
    ensure_backup_dir()
    
    # 获取数据库配置
    try:
        db_settings = get_db_settings()
    except ValueError as e:
        logger.error(f"获取数据库设置失败: {e}")
        return False
    
    # 根据数据库引擎选择备份方法
    engine = db_settings['engine'].lower()
    logger.info(f"检测到数据库引擎: {engine}")
    
    if 'mysql' in engine:
        backup_file = backup_mysql(db_settings)
    elif 'postgresql' in engine:
        backup_file = backup_postgresql(db_settings)
    elif 'sqlite' in engine:
        backup_file = backup_sqlite(db_settings)
    else:
        logger.error(f"不支持的数据库引擎: {engine}")
        return False
    
    # 清理旧备份
    cleanup_old_backups()
    
    if backup_file:
        logger.info("备份过程成功完成")
        return True
    else:
        logger.error("备份过程失败")
        return False


if __name__ == "__main__":
    success = run_backup()
    sys.exit(0 if success else 1) 